from abc import ABC, abstractmethod
import torch
import math
from .seq_learning import spike_seq_to_time


class ErrorFunction(ABC):
    def __init__(self, output_spikes: torch.Tensor, desired_spikes: torch.Tensor):
        self.output_spikes = output_spikes
        self.desired_spikes = desired_spikes
    
    @abstractmethod
    def spike_sequence_error() -> float: pass
    
    @abstractmethod
    def spike_neuron_error() -> torch.Tensor: pass

class CauchySchwarzDistance(ErrorFunction):
    def __init__(self, output_spikes: torch.Tensor, desired_spikes: torch.Tensor):
        super().__init__(output_spikes, desired_spikes)
        self.SIGMA = 5

    def spike_sequence_error(self) -> float:
        return self.cs_distance_error()
    
    def spike_neuron_error(self) -> list:
        return self.cs_distance_neuron_error()

    def cs_distance_error(self) -> float:
        total_error = 0
        for i in range(self.output_spikes.shape[0]):
            sij = self.spike_inner_product(self.output_spikes[i], self.desired_spikes[i])
            sii = self.spike_inner_product(self.output_spikes[i], self.output_spikes[i])
            sjj = self.spike_inner_product(self.desired_spikes[i], self.desired_spikes[i])
            total_error += sij / math.sqrt(sii * sjj)
        return total_error / self.output_spikes.shape[0]

    def cs_distance_neuron_error(self) -> torch.Tensor:
        output_every_error = []
        for i in range(self.output_spikes.shape[0]):
            sij = self.spike_inner_product(self.output_spikes[i], self.desired_spikes[i])
            sii = self.spike_inner_product(self.output_spikes[i], self.output_spikes[i])
            sjj = self.spike_inner_product(self.desired_spikes[i], self.desired_spikes[i])
            output_every_error.append(sij / math.sqrt(sii * sjj))
        output_every_error = torch.Tensor(output_every_error)
        return output_every_error

    def spike_inner_product(self, si: torch.Tensor, sj: torch.Tensor):
        si, sj = spike_seq_to_time(si), spike_seq_to_time(sj)
        # print(si, sj)
        #计算差值矩阵，shape(si.length, sj.length)，m[i][j]=si[i]-sj[j]
        diff_matrix = si.unsqueeze(1) - sj.unsqueeze(0)
        #计算核函数值矩阵
        kernel_values = torch.exp(-(diff_matrix ** 2) / (2 * self.SIGMA * self.SIGMA))
        #内积结果为核函数值矩阵元素之和
        inner_product = torch.sum(kernel_values)
        return inner_product

# if __name__ == "__main__":
    
